package com.gukun.springboot.tdengine.service.backuprecover.recover;

import com.zaxxer.hikari.HikariDataSource;
import org.apache.avro.Schema;
import org.apache.avro.file.DataFileReader;
import org.apache.avro.generic.GenericDatumReader;
import org.apache.avro.generic.GenericRecord;
import org.apache.hadoop.thirdparty.com.google.common.util.concurrent.RateLimiter;

import java.io.File;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.BlockingDeque;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.BiFunction;

/**
 * @author zhang
 */
public class RecoverDbDataTask implements Runnable {

    private final BlockingDeque<File> fileQueue;

    private static final String INSERT_SQL = "insert into %s(ts,v)values(?,?);";
    private static final int BATCH_SIZE = 16000;
    private final HikariDataSource tdEngineConnPool;
    private final String dbName;
    private final AtomicLong  dataNum;

    public RecoverDbDataTask(BlockingDeque<File> fileQueue, String dbName, AtomicLong dataNum, HikariDataSource tdEngineConnPool) {
        this.fileQueue = fileQueue;
        this.tdEngineConnPool = tdEngineConnPool;
        this.dbName = dbName;
        this.dataNum = dataNum;
    }

    @Override
    public void run() {
        List<Long> tsList = new ArrayList<>(BATCH_SIZE);
        List<Object> vList = new ArrayList<>(BATCH_SIZE);
        PreparedStatement pstmt = null;
        // 函数式接口定义
        BiFunction<String, String, Object> tagValueFunction;
        RateLimiter rateLimiter = RateLimiter.create(160000);
        try(Connection connection = tdEngineConnPool.getConnection()){
            try (Statement statement = connection.createStatement()) {
                statement.execute("USE " + dbName);
            }
            while (true) {
                File file = fileQueue.poll(100, TimeUnit.MILLISECONDS);
                if(file == null && fileQueue.isEmpty()){
                    break;
                }
                if(file == null){
                    continue;
                }
                try (DataFileReader<GenericRecord> dataFileReader = new DataFileReader<>(file, new GenericDatumReader<>())) {
                    Schema schema = dataFileReader.getSchema();
                    String tableName = schema.getName();
                    String stableName = schema.getField("stable").defaultVal().toString();
                    tagValueFunction = createTagValueFunction(stableName);
                    String insert = String.format(INSERT_SQL, tableName);
                    GenericRecord record = null;
                    pstmt = connection.prepareStatement(insert);
                    while (dataFileReader.hasNext()) {
                        rateLimiter.acquire();
                        record = dataFileReader.next();
                        tsList.add(Long.valueOf(record.get("ts").toString()));
                        vList.add(tagValueFunction.apply(record.get("v").toString(), stableName));
                        if (tsList.size() >= BATCH_SIZE) {
                            dataNum.getAndAdd(tsList.size());
                            writeData(tsList, vList, pstmt);
                        }
                    }
                    dataNum.getAndAdd(tsList.size());
                    writeData(tsList, vList, pstmt);
                    pstmt.close();
                } finally {
                    if (pstmt != null) {
                        pstmt.close();
                    }
                }
            }
        }catch (Exception e){
            throw new RuntimeException(e);
        }
    }

    private BiFunction<String, String, Object> createTagValueFunction(String stableName) {
        return switch (stableName) {
            case "tint" -> (v, name) -> Integer.parseInt(v);
            case "tfloat", "tdouble" -> (v, name) -> Double.parseDouble(v);
            case "tbool" -> (v, name) -> Boolean.parseBoolean(v);
            default -> (v, name) -> v;
        };
    }


    public void writeData(List<Long> tsList, List<Object> vList, PreparedStatement pstmt) throws SQLException {
        if(tsList.isEmpty() || vList.isEmpty()){
            return;
        }
        for(int i = 0; i < tsList.size(); i++){
            pstmt.setLong(1, tsList.get(i));
            pstmt.setObject(2, vList.get(i));
            pstmt.addBatch();
        }
        pstmt.executeBatch();
        pstmt.clearBatch();
        tsList.clear();
        vList.clear();
    }
}
